# ==============================================================================
# SCRIPT 01: PROCESSAMENTO DE DADOS (2023)
# ==============================================================================

# ------------------------------------------------------------------------------
# 1. OBJETIVO
# ------------------------------------------------------------------------------
# Este script realiza o pré-processamento completo dos dados de 2023, seguindo
# exatamente os mesmos passos do script de 2022 para garantir a comparabilidade.

# ------------------------------------------------------------------------------
# 2. BIBLIOTECAS NECESSÁRIAS
# ------------------------------------------------------------------------------
import pandas as pd
import numpy as np
import warnings

warnings.filterwarnings("ignore")

# ------------------------------------------------------------------------------
# 3. FUNÇÕES AUXILIARES
# ------------------------------------------------------------------------------

def converter_numero_brasileiro(valor):
    """Converte uma string numérica em formato brasileiro para float."""
    if isinstance(valor, str):
        valor = valor.replace(".", "").replace(",", ".")
    try:
        return float(valor)
    except (ValueError, TypeError):
        return np.nan

# ------------------------------------------------------------------------------
# 4. CONFIGURAÇÕES E CONSTANTES
# ------------------------------------------------------------------------------

# Caminhos dos arquivos de entrada
DAIR_PATH = "../2023_data/dados_dair_2023.csv"
ISP_PATH = "../2023_data/dados_isp_2023_atualizado.xlsx"

# Caminho do arquivo de saída
OUTPUT_PATH = "../2023_data/base_final_2023.csv"

# Lista de segmentos de investimento
SEGMENTOS = [
    "Renda Fixa", "Renda Variável", "Investimentos no Exterior",
    "Investimentos Estruturados", "Fundos Imobiliários",
    "Empréstimos Consignados", "Disponibilidades Financeiras"
]

# Lista de erros manifestos (se houver)
ERROS_MANIFESTOS = [] # Nenhum erro manifesto identificado em 2023

# ------------------------------------------------------------------------------
# 5. EXECUÇÃO DO SCRIPT
# ------------------------------------------------------------------------------

print("="*80)
print("INICIANDO PROCESSAMENTO DE DADOS - 2023")
print("="*80)

# --- 5.1. Carregamento e Limpeza dos Dados do DAIR ---
print("\n1. Carregando e limpando dados do DAIR...")

df_dair = pd.read_csv(DAIR_PATH, sep=";", encoding="utf-8-sig", dtype=str)
print(f"  - {len(df_dair)} registros carregados do DAIR.")

df_dair.rename(columns={
    "CNPJ do RPPS": "CNPJ",
    "Mês": "Competência",
    "Segmento": "Segmento",
    "Valor Total Atual": "Valor"
}, inplace=True)

df_dair["Valor"] = df_dair["Valor"].apply(converter_numero_brasileiro)

df_dair = df_dair[df_dair["Segmento"].isin(SEGMENTOS)]
print(f"  - {len(df_dair)} registros após filtrar segmentos.")

df_dair.dropna(subset=["Valor"], inplace=True)
df_dair = df_dair[df_dair["Valor"] >= 0]
print(f"  - {len(df_dair)} registros após remover valores nulos/negativos.")

# --- 5.2. Cálculo dos Percentuais de Alocação ---
print("\n2. Calculando percentuais de alocação...")

df_dair["Valor_Total_RPPS_Mes"] = df_dair.groupby(["CNPJ", "Competência"])["Valor"].transform("sum")
df_dair["Percentual"] = (df_dair["Valor"] / df_dair["Valor_Total_RPPS_Mes"]) * 100
df_dair["Percentual"].fillna(0, inplace=True)

df_dair_dez = df_dair[df_dair["Competência"] == "122023"].copy()
print(f"  - {len(df_dair_dez['CNPJ'].unique())} RPPS encontrados em Dezembro de 2023.")

df_pivot = df_dair_dez.pivot_table(
    index="CNPJ",
    columns="Segmento",
    values="Percentual",
    aggfunc="sum"
).fillna(0)

df_pivot.reset_index(inplace=True)
print(f"  - Base de dados pivotada com {len(df_pivot)} RPPS.")

# --- 5.3. Criar Mapeamento ENTE+UF → CNPJ a partir do DAIR ---
print("\n3. Criando mapeamento ENTE+UF → CNPJ...")

# Extrair ENTE e UF do DAIR
df_mapeamento = df_dair[["CNPJ", "Ente", "UF"]].drop_duplicates()
df_mapeamento["ENTE_PADRAO"] = df_mapeamento["Ente"].str.strip().str.upper()
df_mapeamento["UF_PADRAO"] = df_mapeamento["UF"].str.strip().str.upper()
df_mapeamento["CHAVE"] = df_mapeamento["ENTE_PADRAO"] + "_" + df_mapeamento["UF_PADRAO"]

# Criar dicionário de mapeamento
mapeamento_cnpj = dict(zip(df_mapeamento["CHAVE"], df_mapeamento["CNPJ"]))
print(f"  - {len(mapeamento_cnpj)} mapeamentos ENTE+UF → CNPJ criados.")

# --- 5.4. Carregamento e Limpeza dos Dados do ISP ---
print("\n4. Carregando e limpando dados do ISP...")

df_isp = pd.read_excel(ISP_PATH, dtype=str)
print(f"  - {len(df_isp)} registros carregados do ISP.")

df_isp.rename(columns={
    "CNPJ": "CNPJ_ORIGINAL",
    "ENTE": "ENTE",
    "UF": "UF",
    "REGIÃO": "REGIÃO",
    "INDICADOR DE ACUMULAÇÃO DE RECURSOS\nCLASSIFICAÇÃO": "IAR_CLASSIFICACAO",
    "INDICADOR DE ACUMULAÇÃO DE RECURSOS\nPONTUAÇÃO": "IAR_PONTUACAO"
}, inplace=True)

df_isp["IAR_PONTUACAO"] = pd.to_numeric(df_isp["IAR_PONTUACAO"], errors="coerce")

df_isp.dropna(subset=["ENTE"], inplace=True)

# Extrair ENTE sem o sufixo " - UF" e criar chave de mapeamento
df_isp["ENTE_PADRAO"] = df_isp["ENTE"].str.strip().str.upper().str.replace(r"\s*-\s*[A-Z]{2}$", "", regex=True)
df_isp["UF_PADRAO"] = df_isp["UF"].str.strip().str.upper()
df_isp["CHAVE"] = df_isp["ENTE_PADRAO"] + "_" + df_isp["UF_PADRAO"]

# Preencher CNPJ usando o mapeamento
df_isp["CNPJ"] = df_isp["CHAVE"].map(mapeamento_cnpj)

# Verificar quantos CNPJs foram preenchidos
cnpjs_preenchidos = df_isp["CNPJ"].notna().sum()
print(f"  - {cnpjs_preenchidos} de {len(df_isp)} RPPS tiveram CNPJ preenchido via mapeamento.")

# Remover registros sem CNPJ
df_isp = df_isp[df_isp["CNPJ"].notna()]
print(f"  - {len(df_isp)} RPPS após remover registros sem CNPJ.")

# --- 5.5. Cruzamento das Bases de Dados ---
print("\n5. Cruzando bases de dados DAIR e ISP...")

df_final = pd.merge(df_pivot, df_isp, on="CNPJ", how="inner")
print(f"  - {len(df_final)} RPPS após cruzamento das bases.")

# --- 5.6. Limpeza Final ---
print("\n6. Realizando limpeza final...")

if ERROS_MANIFESTOS:
    df_final["ENTE_PADRAO_CHECK"] = df_final["ENTE_PADRAO"].str.strip().str.upper()
    antes_exclusao = len(df_final)
    df_final = df_final[~df_final["ENTE_PADRAO_CHECK"].str.contains("|".join(ERROS_MANIFESTOS), na=False)]
    depois_exclusao = len(df_final)
    print(f"  - {antes_exclusao - depois_exclusao} RPPS (erros manifestos) excluídos.")
else:
    print("  - Nenhum erro manifesto para excluir.")

print(f"  - Amostra final: {len(df_final)} RPPS.")

# Manter as colunas relevantes
colunas_finais = ["CNPJ", "ENTE", "UF", "REGIÃO", "IAR_CLASSIFICACAO", "IAR_PONTUACAO"] + SEGMENTOS
df_final = df_final[colunas_finais]

# --- 5.7. Salvando a Base Final ---
print(f"\n7. Salvando base de dados final em: {OUTPUT_PATH}")
df_final.to_csv(OUTPUT_PATH, index=False)

print("\n" + "="*80)
print("PROCESSAMENTO DE DADOS DE 2023 CONCLUÍDO COM SUCESSO!")
print("="*80)
